

import sys 
sys.path.append("..") 
sys.path.append("./") 

import numpy as np
import cv2
import h5py
import os
import sys
import torch
import parser_3d
import logging
import sklearn
from os.path import join
from datetime import datetime
from torch.utils.model_zoo import load_url
from google_drive_downloader import GoogleDriveDownloader as gdd

import util

import commons
import datasets.datasets_ws as datasets_ws
from model.search.LocalFeatureSet import LocalFeatureSet



from utils.draw import draw_keypoints

def read_image(impath):
    """ Read image as grayscale and resize to img_size.
    Inputs
        impath: Path to input image.
        img_size: (W, H) tuple specifying resize size.
    Returns
        grayim: float32 numpy array sized H x W with values in range [0, 1].
    """
    grayim = cv2.imread(impath)
    if grayim is None:
        raise Exception('Error reading image %s' % impath)
    # Image is resized via opencv.
    interp = cv2.INTER_AREA
    #grayim = cv2.resize(grayim, (img_size[1], img_size[0]), interpolation=interp)
    grayim = (grayim.astype('float32') / 255.)

    return grayim

img_name = "@0584291.21@4476924.92@17@T@040.43870@-080.00613@004702@11@@@@@@pitch1_yaw12@.jpg"
# feature_path = ""

args = parser_3d.parse_arguments()
# args.resume = "checkpoints/retrievalSfM120k-vgg16-gem-b4dcdc6.pth"
args.dataset_name = "pitts30k"

args.features_dir = "features_MaskANMS_20230602_172723/"

from datasets.dataset_sp import DatasetSP
######################################### DATASETS #########################################

if args.datasets_type == "sp":
    test_ds = DatasetSP(args, args.datasets_folder, args.dataset_name, "test")
else:
    test_ds = datasets_ws.BaseDataset(args, args.datasets_folder, args.dataset_name, "test")
logging.info(f"Test set: {test_ds}")

feature_index = LocalFeatureSet(args)
logging.debug("Read features set")
res = feature_index.read()

imgid = test_ds.get_imgid_by_filename(img_name)
if imgid == -1:
    pass

features = feature_index.get_features_by_id(imgid)

impath = img_name
for path in test_ds.images_paths:
    if impath in path:
        impath = path
        break


img = read_image(impath)
pts = features["keypoints"]

img_pts = draw_keypoints(img, pts)

cv2.imshow('img',img_pts)
cv2.waitKey(0)
